

import os
import xlwings as xw
import numpy as np
import pandas as pd

SRC_FILE = '61wan.xlsx'
DST_FOLDER = './dst'

AREA_DICT = {
    '玉溪市': '530400',
    '市辖区': '530401',
    '红塔区': '530402',
    '江川县': '530421',
    '澄江县': '530422',
    '通海县': '530423',
    '华宁县': '530424',
    '易门县': '530425',
    '峨山县': '530426',
    '新平县': '530427',
    '元江县': '530428',
}

AREA_DICT1 = {
    '玉溪市和红塔区': ['530400', '530402'],
    '市辖区': '530401',
    '江川县': '530421',
    '澄江县': '530422',
    '通海县': '530423',
    '华宁县': '530424',
    '易门县': '530425',
    '峨山县': '530426',
    '新平县': '530427',
    '元江县': '530428',
}


def read_from_xlsx(filepath):
    wb = xw.Book(filepath)
    sht = wb.sheets[0]

    info = sht.used_range
    nrows = info.last_cell.row
    ncols = info.last_cell.column

    df = sht.range((1, 1), (nrows, ncols)).options(
        pd.DataFrame, index=False, header=False).value
    df.columns = list('ABCDE')
    wb.close()
    return df


def write_to_workbook(data):
    name, df = data

    df = df.drop(['F'], axis=1)
    df['C'].astype(str)

    wb = xw.Book()
    sht = wb.sheets[0]
    head = ['序号', '受评主体名称', '统一社会信用代码', '综合得分', '评价等级']

    sht.range((1, 1), (1, 5)).value = head
    sht.range((1, 2)).column_width = 40
    sht.range((1, 3)).column_width = 20

    sht.range('C:C').number_format = '@'
    sht.range('A2').options(index=False, header=False).value = df

    wb.save(DST_FOLDER + '/' + name + '（' + str(df.shape[0]) + '家）' + '.xlsx')
    wb.close()


def check_code(df, code):
    if code != None:
        if isinstance(code, list):
            return df.isin(code)
        else:
            return df == code
    else:
        return False


def main():

    if not os.path.exists(DST_FOLDER):
        os.mkdir(DST_FOLDER)

    try:
        app = xw.App(visible=False)

        df = read_from_xlsx(SRC_FILE)
        print('Read data done')

        # 提取区域代码为新一列
        df['F'] = df['C'].apply(lambda x: x[2:8] if x != None else None)

        for area, code in AREA_DICT1.items():
            print(f'Process {area}')
            data = (area, df[check_code(df['F'], code)])
            write_to_workbook(data)
        print('Done')
    finally:
        app.quit()


if __name__ == '__main__':
    main()
